/**
* EPOLLONESHOT demo for epoll
* @date 2021/07/18
* @author Martin Zhang(zhangming_233@163.com)
*/
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>


#define MAX_EVENT_NUMBER 1024
#define BUFFER_SIZE      1024

typedef struct sockaddr SA;

struct fds {
    int epollfd;
    int sockfd;
};

int setnonblocking(int fd)
{
    int old_option = fcntl(fd, F_GETFL);
    int new_option = old_option | O_NONBLOCK;
    fcntl(fd, F_SETFL, new_option);
    return old_option;
}

/* 将fd上的EPOLLIN和EPOLLET事件注册到epollfd指示的epoll内核事件表中，参数oneshot指定是否注册fd上的EPOLLONESHOT事件  */
void addfd (int epollfd, int fd, bool oneshot)
{
    struct epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLIN | EPOLLET;
    if (oneshot) 
        event.events |= EPOLLONESHOT;

    epoll_ctl(epollfd, EPOLL_CTL_ADD, fd, &event);
    setnonblocking(fd);
}

/* 重置fd上的事件。这样操作之后，尽管fd上的EPOLLONESHOT事件被注册，但是OS仍然会触发fd上的EPOLLIN事件，且只触发一次 */
void reset_oneshot(int epollfd, int fd)
{
    struct epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLIN | EPOLLET | EPOLLONESHOT;
    epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, &event);
}

/* 工作线程 */
void *worker(void *arg)
{
    int sockfd = ((struct fds *)arg)->sockfd;
    int epollfd = ((struct fds *)arg)->epollfd;

    printf("start new thread to receive data on fd: %d\n", sockfd);
    char buf[BUFFER_SIZE];
    memset(buf, '\0', BUFFER_SIZE);

    /* 循环读取sockfd上的数据，直到遇到EAGAIN错误 */
    while (1) {
        int ret = recv(sockfd, buf, BUFFER_SIZE - 1, 0);
        if (ret == 0) {
            close(sockfd);
            printf("foreiner closed the connection\n");
            break;
        }
        else if (ret < 0) {
            if (errno == EAGAIN) { // 在ET模式下, 对non-blocking fd调用阻塞操作recv, 该操作可能没有完成就返回, 从而产生该错误码, 但不会破坏socket
                reset_oneshot(epollfd, sockfd); // 重置sockfd上注册的EPOLLONESHOT事件以及其他事件
                printf("read later\n");
                break;
            }
        }
        else {
            printf("get content: %s\n", buf);
            /* 休眠5秒，模拟数据处理过程 */
            sleep(5);
        }
    }

    printf("end thread receiving data on fd: %d\n", sockfd);
}


/**
 * epoll EPOLLONESHOT例程
 * @note build way(just like): 
 * gcc epoll_EPOLLONESHOT.c -pthread -std=gnu99 -o epolloneshot
 */
int main(int argc, char *argv[])
{
    if (argc <= 2) {
        printf("usage: %s ip_address port_number\n", argv[0]);
        return 1;
    }

    const char *ip = argv[1];
    int port = atoi(argv[2]);

    int ret = 0;
    struct sockaddr_in address;
    bzero(&address, sizeof(address));
    address.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &address.sin_addr);
    address.sin_port = htons(port);

    int listenfd = socket(AF_INET, SOCK_STREAM, 0);
    assert(listenfd >= 0);

    ret = bind(listenfd, (SA *)&address, sizeof(address));
    assert(ret != -1);

    ret = listen(listenfd, 5);
    assert(ret != -1);

    struct  epoll_event events[MAX_EVENT_NUMBER];
    int epollfd = epoll_create(5); /* 5没有意义，但必须>0 */
    assert(epollfd != -1);
    /* 注意：监听socket listenfd 上是不能注册EPOLLONESHOT事件的，否则应用程序只能处理一个客户连接！因为后续的客户连接请求将不再触发listenfd上的EPOLLIN事件 */
    addfd(epollfd, listenfd, false);

    while (1) {
        int ret = epoll_wait(epollfd, events, MAX_EVENT_NUMBER, -1);
        if (ret < 0) {
            printf("epoll failure\n");
            break;
        }

        for (int i = 0; i < ret; i++) {
            int sockfd = events[i].data.fd;
            if (sockfd == listenfd) {
                struct sockaddr_in client_address;
                socklen_t client_addrlength = sizeof(client_address);
                int connfd = accept(listenfd, (SA *)&client_address, &client_addrlength);
                /* 对每个非监听文件描述符都注册EPOLLONESHOT事件 */
                addfd(epollfd, connfd, true);
            }
            else if (events[i].events & EPOLLIN ) { /* 普通可读数据就绪 */
                pthread_t thread;
                struct fds fds_for_new_worker;
                fds_for_new_worker.epollfd = epollfd;
                fds_for_new_worker.sockfd = sockfd;

                /* 新启动一个工作线程为sockfd服务 */
                pthread_create(&thread, NULL, worker, (void *)&fds_for_new_worker);
            }
            else {
                printf("something else happend\n");
            }
        }
    }

    close(listenfd);
    return 0;
}
